package com.gitee.wsl.io.channel.read.ext


import com.gitee.wsl.data.pool.sample.SimpleCreatePool
import com.gitee.wsl.io.channel.ByteChannel
import com.gitee.wsl.io.channel.read.ByteReadChannel
import com.gitee.wsl.io.channel.write.ext.writeFully
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.launch


private const val DEFAULT_POOL_ARRAY_SIZE = 4096
private const val DEFAULT_POOL_CAPACITY = 128

val ByteArrayPool: SimpleCreatePool<ByteArray> = object : SimpleCreatePool<ByteArray>(
    DEFAULT_POOL_CAPACITY
) {
    override fun produceInstance(): ByteArray = ByteArray(DEFAULT_POOL_ARRAY_SIZE)
}

/**
 * Split source [ByteReadChannel] into 2 new ones.
 * Cancel of one channel in split (input or both outputs) cancels other channels.
 */
fun ByteReadChannel.split(coroutineScope: CoroutineScope): Pair<ByteReadChannel, ByteReadChannel> {
    val first = ByteChannel(autoFlush = true)
    val second = ByteChannel(autoFlush = true)

    coroutineScope.launch {
        val buffer = ByteArrayPool.acquire()
        try {
            while (!isClosedForRead) {
                val read = this@split.readAvailable(buffer)
                if (read <= 0) continue
                listOf(
                    async { first.writeFully(buffer, 0, read) },
                    async { second.writeFully(buffer, 0, read) }
                ).awaitAll()
            }

            closedCause?.let { throw it }
        } catch (cause: Throwable) {
            this@split.cancel(cause)
            first.cancel(cause)
            second.cancel(cause)
        } finally {
            ByteArrayPool.release(buffer)
            first.close()
            second.close()
        }
    }.invokeOnCompletion {
        it ?: return@invokeOnCompletion
        first.cancel(it)
        second.cancel(it)
    }

    return first to second
}